
# 郑州大学科研平台专用

import os
import random 
random.seed(0)

# 在终端控制台输入指令：pwd 查看当前所在目录
# 修改 utils/datasets.py 文件的349行，将 labels 改为 yolo_labels

# yolo标签路径
label_path = '/home/jovyan/exp_2715/dataset-445/yolo_labels'

# 划分后txt文件保存路径
save_path = '/home/jovyan/exp_2715/data/my_dataset/'

# yolo数据集图片路径
image_path = '/home/jovyan/exp_2715/dataset-445/images/'

# 数据集总数量为1
trainval_percent = 1

# 划分比例
train_percent = 0.8
 
temp = os.listdir(label_path)

num=len(temp)  

list=range(num)  

tv=int(num*trainval_percent)

tr=int(tv*train_percent) 

trainval= random.sample(list,tv)  

train=random.sample(trainval,tr)  

 
print("数据集总数量",tv)
print("训练集数量",tr)

ftrainval = open(os.path.join(save_path,'trainval.txt'), 'w')  
ftest = open(os.path.join(save_path,'test.txt'), 'w')  
ftrain = open(os.path.join(save_path,'train.txt'), 'w')  
fval = open(os.path.join(save_path,'valid.txt'), 'w')  

# 根据数据集图片名称自己改png还是jpg
for i  in list:  
    name=temp[i][:-4] 
    if i in trainval:  
        ftrainval.write(image_path+name+'.png'+'\n' )  
        if i in train:  
            ftrain.write(image_path+name+'.png'+'\n' )  
        else:  
            fval.write(image_path+name+'.png'+'\n' )  
    else:  
        ftest.write(image_path+name+'.png'+'\n' )  

ftrainval.close()  
ftrain.close()  
fval.close()  
ftest .close()

